#!/bin/env python

import os
import readline
import subprocess

HADOOP_EXECUTABLE = '/home/kelvin/hadoop/bin/hadoop'
HDFS_HOME = '/user/kelvin'

def travel(pwd, d):
    if d == '' or d == '.':
        pass
    elif d == '..':
        i = pwd.rfind('/')
        pwd = pwd[:i]
    elif pwd[-1] != '/':
        pwd += '/' + d
    else:
        pwd += d
    return pwd

def chdir(pwd, path=None):

    if path is None:
        return HDFS_HOME

    if path.startswith('/'):
        pwd = '/'
        path = path[1:]

    ds = path.split('/')
    for d in ds:
        pwd = travel(pwd, d)

    return pwd

def retrieve_items(output):
    return [line.split(' ')[-1].split('/')[-1] for line in output.split('\n') if line.split(' ')[-1].startswith('/')]

pwd = HDFS_HOME
items = None

def complete(prefix, index):
    global pwd
    global items

    if items is None:
        output = subprocess.check_output([HADOOP_EXECUTABLE, 'dfs', '-ls', pwd])
        items = retrieve_items(output)

    ws = [w for w in items if w.startswith(prefix)]
    try:
        return ws[index]
    except IndexError:
        return None

def main():
    global pwd
    global items

    while True:
        try:
            cmd = raw_input('[hdfs-shell %s]$ ' % pwd).strip()
        except EOFError:
            print 'exit'
            break

        tokens = cmd.split(' ')
        if tokens[0] == '':
            pass
        elif tokens[0] == 'quit' or tokens[0] == 'exit':
            break
        elif tokens[0] == 'cd':
            pwd_old = pwd
            if len(tokens) > 1:
                pwd = chdir(pwd, tokens[1])
            else:
                pwd = chdir(pwd)
            if pwd != pwd_old:
                items = None
        elif tokens[0] == 'pwd':
            print pwd
        elif tokens[0] in ['cat', 'text', 'ls', 'lsr', 'rm', 'rmr', 'du', 'dus', 'count', 'mkdir', 'stat', 'tail', 'mv', 'cp']:
            if tokens[0] == 'ls' and len(tokens) == 1:
                output = subprocess.check_output([HADOOP_EXECUTABLE, 'dfs', '-ls', pwd])
                items = retrieve_items(output)
                if output:
                    print output,
            else:
                if len(tokens) == 1: tokens.append('')
                for i in range(1, len(tokens)):
                    if tokens[i] in ['<', '<<', '>', '>>', '|']: break
                    if not tokens[i].startswith('-'): 
                        tokens[i] = chdir(pwd, tokens[i])
                cmd = ' '.join(tokens)
                os.system('%s dfs -%s' % (HADOOP_EXECUTABLE, cmd))
        else:
            os.system('%s dfs -%s' % (HADOOP_EXECUTABLE, cmd))

if __name__ == '__main__':
    readline.parse_and_bind('tab: complete')
    readline.set_completer(complete)
    main()
